use std::{
    fmt::Display,
    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
};

use crate::lexer::{HasLocation, Location};

static INDENT: AtomicUsize = AtomicUsize::new(0);
static SIMPLE: AtomicBool = AtomicBool::new(false);

pub struct Module {
    pub functions: Vec<Function>,
}
impl Display for Module {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);

        writeln!(f, "{}Module:", " ".repeat(indent))?;
        INDENT.fetch_add(2, Ordering::Relaxed);
        for function in &self.functions {
            write!(f, "{}", function)?;
        }
        INDENT.fetch_sub(2, Ordering::Relaxed);
        Ok(())
    }
}

pub struct Function {
    pub proto: Prototype,
    pub body: ExprList,
}
impl Display for Function {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);

        writeln!(f, "{}Function", " ".repeat(indent))?;
        INDENT.fetch_add(2, Ordering::Relaxed);
        write!(f, "{}", self.proto)?;
        writeln!(
            f,
            "{}Params: [{}]",
            " ".repeat(indent + 2),
            self.proto
                .args
                .iter()
                .map(|a| a.name.clone())
                .collect::<Vec<String>>()
                .join(", ")
        )?;
        write!(f, "{}", self.body)?;
        INDENT.fetch_sub(2, Ordering::Relaxed);
        Ok(())
    }
}

pub struct Prototype {
    pub name: String,
    pub args: Vec<VariableExpr>,
    pub(crate) location: Location,
}
impl Display for Prototype {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);

        writeln!(
            f,
            "{}Proto '{}' {}",
            " ".repeat(indent),
            self.name.clone(),
            self.location()
        )
    }
}
impl HasLocation for Prototype {
    fn location(&self) -> Location {
        self.location.clone()
    }
}

pub struct ExprList {
    pub exprs: Vec<Expr>,
}
impl Display for ExprList {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);

        writeln!(f, "{}Block {{", " ".repeat(indent))?;
        INDENT.fetch_add(2, Ordering::Relaxed);
        for expr in &self.exprs {
            write!(f, "{}", expr)?;
        }
        INDENT.fetch_sub(2, Ordering::Relaxed);
        writeln!(f, "{}}}  // Block", " ".repeat(indent))?;
        Ok(())
    }
}

pub enum Expr {
    NumberExpr(NumberExpr),
    LiteralExpr(LiteralExpr),
    VariableExpr(VariableExpr),
    VarDeclExpr(VarDeclExpr),
    ReturnExpr(ReturnExpr),
    BinaryExpr(BinaryExpr),
    CallExpr(CallExpr),
    PrintExpr(PrintExpr),
}
impl Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::NumberExpr(number_expr) => write!(f, "{}", number_expr),
            Expr::LiteralExpr(literal_expr) => write!(f, "{}", literal_expr),
            Expr::VariableExpr(variable_expr) => write!(f, "{}", variable_expr),
            Expr::VarDeclExpr(var_decl_expr) => write!(f, "{}", var_decl_expr),
            Expr::ReturnExpr(return_expr) => write!(f, "{}", return_expr),
            Expr::BinaryExpr(binary_expr) => write!(f, "{}", binary_expr),
            Expr::CallExpr(call_expr) => write!(f, "{}", call_expr),
            Expr::PrintExpr(print_expr) => write!(f, "{}", print_expr),
        }
    }
}
impl HasLocation for Expr {
    fn location(&self) -> Location {
        match self {
            Expr::NumberExpr(_number_expr) => todo!(),
            Expr::LiteralExpr(literal_expr) => literal_expr.location(),
            Expr::VariableExpr(variable_expr) => variable_expr.location(),
            Expr::VarDeclExpr(var_decl_expr) => var_decl_expr.location(),
            Expr::ReturnExpr(_return_expr) => todo!(),
            Expr::BinaryExpr(binary_expr) => binary_expr.location(),
            Expr::CallExpr(call_expr) => call_expr.location(),
            Expr::PrintExpr(print_expr) => print_expr.location(),
        }
    }
}

pub struct PrintExpr {
    pub arg: Box<Expr>,
    pub(crate) location: Location,
}
impl Display for PrintExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);
        writeln!(f, "{}Print [ {}", " ".repeat(indent), self.location())?;
        INDENT.fetch_add(2, Ordering::Relaxed);
        write!(f, "{}", self.arg)?;
        INDENT.fetch_sub(2, Ordering::Relaxed);
        writeln!(f, "{}]", " ".repeat(indent))?;
        Ok(())
    }
}
impl HasLocation for PrintExpr {
    fn location(&self) -> Location {
        self.location.clone()
    }
}

pub struct CallExpr {
    pub callee: String,
    pub args: Vec<Expr>,
    pub(crate) location: Location,
}
impl Display for CallExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);
        writeln!(
            f,
            "{}Call '{}' [ {}",
            " ".repeat(indent),
            self.callee.clone(),
            self.location()
        )?;
        INDENT.fetch_add(2, Ordering::Relaxed);
        for arg in &self.args {
            write!(f, "{}", arg)?;
        }
        INDENT.fetch_sub(2, Ordering::Relaxed);
        writeln!(f, "{}]", " ".repeat(indent))?;
        Ok(())
    }
}
impl HasLocation for CallExpr {
    fn location(&self) -> Location {
        self.location.clone()
    }
}

pub struct BinaryExpr {
    pub lhs: Box<Expr>,
    pub op: BinaryOperator,
    pub rhs: Box<Expr>,
}
impl Display for BinaryExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);

        writeln!(
            f,
            "{}BinOp: {} {}",
            " ".repeat(indent),
            self.op,
            self.location()
        )?;
        INDENT.fetch_add(2, Ordering::Relaxed);
        write!(f, "{}", &*self.lhs)?;
        write!(f, "{}", &*self.rhs)?;
        INDENT.fetch_sub(2, Ordering::Relaxed);
        Ok(())
    }
}
impl HasLocation for BinaryExpr {
    fn location(&self) -> Location {
        self.rhs.location()
    }
}

pub enum BinaryOperator {
    Add,
    Mul,
}
impl Display for BinaryOperator {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BinaryOperator::Add => write!(f, "+"),
            BinaryOperator::Mul => write!(f, "*"),
        }
    }
}

pub struct ReturnExpr {
    pub expr: Option<Box<Expr>>,
}
impl Display for ReturnExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);

        writeln!(f, "{}Return", " ".repeat(indent))?;
        if let Some(expr) = &self.expr {
            INDENT.fetch_add(2, Ordering::Relaxed);
            write!(f, "{}", &*expr)?;
            INDENT.fetch_sub(2, Ordering::Relaxed);
        }
        Ok(())
    }
}

pub struct VarDeclExpr {
    pub name: String,
    pub r#type: VarType,
    pub init_val: Box<Expr>,
    pub(crate) location: Location,
}
impl Display for VarDeclExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);

        writeln!(
            f,
            "{}VarDecl {}{} {}",
            " ".repeat(indent),
            self.name.clone(),
            self.r#type,
            self.location()
        )?;
        INDENT.fetch_add(2, Ordering::Relaxed);
        write!(f, "{}", &*self.init_val)?;
        INDENT.fetch_sub(2, Ordering::Relaxed);
        Ok(())
    }
}
impl HasLocation for VarDeclExpr {
    fn location(&self) -> Location {
        self.location.clone()
    }
}

#[derive(Default)]
pub struct VarType {
    pub shape: Vec<u64>,
}
impl Display for VarType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "<{}>",
            self.shape
                .iter()
                .map(|&s| format!("{}", s))
                .collect::<Vec<String>>()
                .join(", ")
        )
    }
}

pub struct VariableExpr {
    pub name: String,
    pub(crate) location: Location,
}
impl Display for VariableExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);
        writeln!(
            f,
            "{}var: {} {}",
            " ".repeat(indent),
            self.name.clone(),
            self.location()
        )
    }
}
impl HasLocation for VariableExpr {
    fn location(&self) -> Location {
        self.location.clone()
    }
}

pub struct LiteralExpr {
    pub values: Vec<Expr>,
    pub dims: Vec<u64>,
    pub(crate) location: Location,
}
impl Display for LiteralExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);
        let simple = SIMPLE.load(Ordering::Relaxed);

        if simple {
            write!(
                f,
                "<{}>[ {}]",
                self.dims
                    .iter()
                    .map(|&d| format!("{}", d))
                    .collect::<Vec<String>>()
                    .join(", "),
                {
                    let ret = self
                        .values
                        .iter()
                        .map(|v| format!("{}", v))
                        .collect::<Vec<String>>()
                        .join(", ");
                    ret
                }
            )?;
        } else {
            SIMPLE.store(true, Ordering::Relaxed);
            writeln!(
                f,
                "{}Literal: <{}>[ {}] {}",
                " ".repeat(indent),
                self.dims
                    .iter()
                    .map(|&d| format!("{}", d))
                    .collect::<Vec<String>>()
                    .join(", "),
                {
                    let ret = self
                        .values
                        .iter()
                        .map(|v| format!("{}", v))
                        .collect::<Vec<String>>()
                        .join(", ");
                    ret
                },
                self.location()
            )?;
            SIMPLE.store(false, Ordering::Relaxed);
        }

        Ok(())
    }
}
impl HasLocation for LiteralExpr {
    fn location(&self) -> Location {
        self.location.clone()
    }
}

pub struct NumberExpr {
    pub val: f64,
    pub(crate) location: Location,
}
impl Display for NumberExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let indent = INDENT.load(Ordering::Relaxed);
        let simple = SIMPLE.load(Ordering::Relaxed);

        if simple {
            write!(f, "{:.6e}", self.val)
        } else {
            writeln!(
                f,
                "{}{:.6e} {}",
                " ".repeat(indent),
                self.val,
                self.location()
            )
        }
    }
}
impl HasLocation for NumberExpr {
    fn location(&self) -> Location {
        self.location.clone()
    }
}
